﻿using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using JasLib.Extensions;

namespace JasLib.Threading
{
    public class AsyncSemaphore : AsyncQueuedSynchronizationPrimitive
    {
        private volatile int _iMaxOccupants = 1;
        public int MaxOccupants
        {
            get { return _iMaxOccupants; }
            set { _iMaxOccupants = value; }
        }

        private readonly bool _bAllowTaskRecursion = false;
        private Queue<TaskCompletionSource<object>> _waitQueue = new Queue<TaskCompletionSource<object>>();

        private readonly Dictionary<int, OccupantInfo> _recursibleOccupants = null;
        private int _iNonRecursableOccupantCount = 0;

        public AsyncSemaphore(int iStartingMaxOccupants = 1, bool bAllowTaskRecursion = false)
            : base(bAllowTaskRecursion)
        {
            _iMaxOccupants = iStartingMaxOccupants;
            _bAllowTaskRecursion = bAllowTaskRecursion;

            if (_bAllowTaskRecursion)
                _recursibleOccupants = new Dictionary<int, OccupantInfo>();
            return;
        }

        protected override void Dispose(bool bFinalizer)
        {
            base.Dispose(bFinalizer);

            var stolenWaitQueue = Interlocked.Exchange(ref _waitQueue, null);
            if (stolenWaitQueue != null)
            {
                var ex = new ObjectDisposedException("SemaphoreGate");
                while (stolenWaitQueue.Count > 0)
                    stolenWaitQueue.Dequeue().SetException(ex);
            }

            return;
        }

        private async void OnTokenCancellation(object objState)
        {
            Queue<TaskCompletionSource<object>> stolenWaitQueue = null;

            using (_objLock.Lock())
            {
                /// This is internal code; don't throw an exception here.
                if (this.Disposed)
                    return;

                await AsyncThreadSwitcher.UnsafeSwitchToThreadPoolAsync();

                GeneralUtilities.Swap(ref stolenWaitQueue, ref _waitQueue);

                /// Release everything. Cancellation means the semaphore is all but disposed.
                Dispose();
            }

            if (stolenWaitQueue != null)
            {
                var ex = new OperationCanceledException((CancellationToken)objState);
                while (stolenWaitQueue.Count > 0)
                    stolenWaitQueue.Dequeue().SetException(ex);
            }
            return;
        }

        public override async Task WaitAsync(TimeSpan? waitTime = null)
        {
            EnforceNotDisposed();
            int iCurrentTaskId = GetCurrentTaskId();

            TaskCompletionSource<object> tcs = null;

            using (_objLock.Lock())
            {
                EnforceNotDisposed();

                if (_waitQueue != null)
                {
                    int iOccupantCount = _bAllowTaskRecursion ? _recursibleOccupants.Count : _iNonRecursableOccupantCount;
                    if (iOccupantCount >= _iMaxOccupants)
                    {
                        tcs = new TaskCompletionSource<object>();
                        _waitQueue.Enqueue(tcs);
                    }
                }
            }

            /// If a token cancellation or dispose happens during this await, it will throw an exception here.
            if (tcs != null)
                await tcs.Task;

            if (_bAllowTaskRecursion)
            {
                using (_objLock.Lock())
                {
                    EnforceNotDisposed();

                    OccupantInfo occupant;
                    if (!_recursibleOccupants.TryGetValue(iCurrentTaskId, out occupant))
                    {
                        occupant = new OccupantInfo();
                        _recursibleOccupants.Add(iCurrentTaskId, occupant);
                    }
                    occupant._iRecursionCount++;
                }
            }
            else
            {
                Interlocked.Increment(ref _iNonRecursableOccupantCount);
            }

            return;
        }

        public void Leave()
        {
            EnforceNotDisposed();
            int iCurrentTaskId = GetCurrentTaskId();

            TaskCompletionSource<object> tcs = null;

            using (_objLock.Lock())
            {
                EnforceNotDisposed();

                if (_bAllowTaskRecursion)
                {
                    OccupantInfo occupant;
                    if (!_recursibleOccupants.TryGetValue(iCurrentTaskId, out occupant))
                        throw new InvalidOperationException("The calling task did not first acquire the semaphore before attempting to release it.");

                    occupant._iRecursionCount--;

                    if (occupant._iRecursionCount == 0)
                    {
                        _recursibleOccupants.Remove(iCurrentTaskId);
                    }
                    else if (occupant._iRecursionCount < 0)
                    {
                        Dispose();
                        throw new InvalidOperationException("The number of releases exceeds the number of acquires. The semaphore should be considered corrupt and will be disposed.");
                    }
                }
                else
                {
                    if (Interlocked.Decrement(ref _iNonRecursableOccupantCount) < 0)
                    {
                        Dispose();
                        throw new InvalidOperationException("The number of releases exceeds the number of acquires. The semaphore should be considered corrupt and will be disposed.");
                    }
                }

                /// Release the next one to take the lock.
                if (!_waitQueue.IsNullOrEmpty())
                    tcs = _waitQueue.Dequeue();
            }

            if (tcs != null)
                tcs.SetResult(null);
            return;
        }
    }
}
